
# KrylovKit version
function mastereq_dt(
    rho::Matrix{ComplexF64},
    params::Tuple{SparseMatrixCSC{ComplexF64, Int64}, SparseMatrixCSC{ComplexF64, Int64}, Vector{SparseMatrixCSC{Float64, Int64}}, Vector{SparseMatrixCSC{Float64, Int64}}, Matrix{ComplexF64}, Matrix{ComplexF64}},
)

    hamiltonian_nh, hamiltonian_nh_dag, lindblad_arr, lindblad_dag_arr, dummy1, dummy2 = params;

    @assert all(isfinite, rho)

    mul!(dummy1, hamiltonian_nh, rho)
    mul!(dummy2, rho, hamiltonian_nh_dag)
    drho = (-im) .* (dummy1 .- dummy2)

    for (lindblad, lindblad_dag) in Iterators.zip(lindblad_arr, lindblad_dag_arr)
        mul!(dummy1, rho, lindblad_dag)
        mul!(dummy2, lindblad, dummy1)

        drho .+= dummy2
    end

    @assert all(isfinite, drho)
    drho
end


function mastereq_dt(
    rho_arr::Vector{Matrix{ComplexF64}},
    params_arr::Vector{Tuple{SparseMatrixCSC{ComplexF64, Int64}, SparseMatrixCSC{ComplexF64, Int64}, Vector{SparseMatrixCSC{Float64, Int64}}, Vector{SparseMatrixCSC{Float64, Int64}}, Matrix{ComplexF64}, Matrix{ComplexF64}}},
)
    drho_arr = Vector{Matrix{ComplexF64}}()

    for (rho, params) in Iterators.zip(rho_arr, params_arr)
        hamiltonian_nh, hamiltonian_nh_dag, lindblad_arr, lindblad_dag_arr, dummy1, dummy2 = params;

        @assert all(isfinite, rho)

        mul!(dummy1, hamiltonian_nh, rho)
        mul!(dummy2, rho, hamiltonian_nh_dag)
        drho = (-im) .* (dummy1 .- dummy2)

        for (lindblad, lindblad_dag) in Iterators.zip(lindblad_arr, lindblad_dag_arr)
            mul!(dummy1, rho, lindblad_dag)
            mul!(dummy2, lindblad, dummy1)

            drho .+= dummy2
        end

        @assert all(isfinite, drho)
        push!(drho_arr, drho)
    end

    drho_arr
end


function energy_subspaces(energies::Matrix{Float64}, max_energy_difference::Float64)

    energy_order = sortperm(energies[:])
    idx_mat = collect(Iterators.product(axes(energies)...))
    sub_manifold_arr::Vector{Vector{Tuple{Int64,Int64}}} = [[idx_mat[energy_order[1]]]]
        
    for ii in energy_order
        if ii == energy_order[1]
            continue
        end

        energy = energies[ii]
        if energy < energies[sub_manifold_arr[end][end]...] + max_energy_difference 
            push!(sub_manifold_arr[end], idx_mat[ii])
        else
            push!(sub_manifold_arr, [idx_mat[ii]])
        end
    end

    sub_manifold_arr
end


function find_energy_subspace_idx(energy_submanifold_arr::Vector{Vector{Tuple{Int64,Int64}}}, idxs::Tuple{Int64,Int64})
    for (ii, energy_submanifold) in enumerate(energy_submanifold_arr)
        if idxs in energy_submanifold
            return (ii, findall(x -> x==idxs, energy_submanifold)[1])
        end
    end
end
